Google JAX is a machine learning framework for transforming numerical functions. It is described as bringing together a modified version of autograd (automatic obtaining of the gradient function through differentiation of a function) and TensorFlow's XLA (Accelerated Linear Algebra). It is designed to follow the structure and workflow of NumPy as closely as possible and works with various existing frameworks such as TensorFlow and PyTorch. The primary functions of JAX are: grad: automatic differentiation jit: compilation vmap: auto-vectorization pmap: Single program, multiple data (SPMD) programming